#include <book.h>

#if HAVE_CONFIG_H
# include <config.h>
#endif

#include <errors.h>
#include <game.h>
#include <movegen.h>
#include <pgn.h>
#include <position.h>
#include <ttable.h>

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>

#define BOOK_MAGIC UINT64_C(0x6b6f6f426b6f6f42)
#define INIT_MHASH_SIZE (1 << 16)
#define MHASH_LOAD_FACTOR 50  // percentage
#define BOOK_COUNT_BITS 16
#define BOOK_MAX_COUNT ((1 << BOOK_COUNT_BITS) - 1)
#define BOOK_GET_KEY(x) ((x) & ~((uint64_t)BOOK_MAX_COUNT))
#define BOOK_GET_COUNT(x) ((x) & ((uint64_t)BOOK_MAX_COUNT))
#define BOOK_ENTRY(h, c) (BOOK_GET_KEY(h) | ((uint64_t)(c)))

/*
 * A simpe hash table with quadratic collision resolution.
 * - bits 0-15 position count
 * - bits 16-63 upper 48 bits of the hash key
 */
static hash_key_t* book_hash = 0;
static uint64_t hash_size = 0;


// Used to build the book.
struct mhash_entry_t
{
  hash_key_t key;
  uint64_t count;
};


static void* encode_uint64(uint64_t x, uint8_t* buf)
{
  for (int i = 0; i < 8; i++) {
    buf[i] = x & 255;
    x >>= 8;
  }
  return buf;
}


static uint64_t decode_uint64(const uint8_t* buf)
{
  uint64_t x = 0;
  for (int i = 0; i < 8; i++)
    x |= ((uint64_t)buf[i]) << (i * 8);
  return x;
}


static int mhash_resize(struct mhash_entry_t** ht,
			size_t old,
			size_t req)
{
  assert(ht != 0);
  assert(req >= old);
  assert(*ht != 0 || old == 0);
  assert((old & (old - 1)) == 0);
  assert((req & (req - 1)) == 0);
  
  if (old == req)
    return 1;

  void* m = malloc(sizeof(struct mhash_entry_t) * req);
  if (m == 0) {
    error_set(ERR_MEMORY);
    return 0;
  }
  
  struct mhash_entry_t* newht = (struct mhash_entry_t*)m;
  for (size_t i = 0; i < req; i++) {
    newht[i].key = 0;
    newht[i].count = 0;
  }

  hash_key_t mask = req - 1;
  struct mhash_entry_t* oldht = *ht;
  for (size_t i = 0; i < old; i++) {
    if (oldht[i].count == 0)
      continue;
    hash_key_t k = oldht[i].key & mask;
    hash_key_t inc = 1;
    while (newht[k].count > 0) {
      k = (k + inc) & mask;
      ++inc;
    }
    newht[k] = oldht[i];
  }

  free(oldht);
  *ht = newht;
  return 1;
}


static int mhash_insert(struct mhash_entry_t** ht,
			 size_t* sz,
			 size_t* occ,
			 hash_key_t key)
{
  if (((*sz) * MHASH_LOAD_FACTOR) <= (*occ) * 100) {
    size_t new_size = ((*sz) == 0 ? INIT_MHASH_SIZE : 2 * (*sz));
    if (!mhash_resize(ht, *sz, new_size))
      return 0;
    *sz = new_size;
  }

  hash_key_t mask = (*sz) - 1;
  hash_key_t k = key & mask;
  struct mhash_entry_t* h = *ht;
  hash_key_t inc = 1;
  while (h[k].count > 0 && h[k].key != key) {
    k = (k + inc) & mask;
    ++inc;
  }
  if (h[k].key == key) {
    ++h[k].count;
  } else {
    h[k].key = key;
    h[k].count = 1;
    ++(*occ);
  }
  return 1;
}


static int insert_pgn_game(FILE* pgn,
			   struct mhash_entry_t** ht,
			   size_t* sz,
			   size_t* occ,
			   int plies)
{
  assert(pgn != 0);
  assert(ht != 0);
  assert(sz != 0);
  assert(*sz == 0 || *ht != 0);
  assert(!(*sz & (*sz - 1)));  // 0, or power of two
  assert(occ != 0);
  assert(plies >= 0);
  
  game_t g;
  if (!game_init(&g))
    return -1;
  int ret = pgn_read(pgn, &g);
  if (ret < 1) {
    game_cleanup(&g);
    return ret;
  }
  
  int len = game_length(&g);
  while (len >= 0) {
    if (len <= plies) {
      const position_t* pos = game_cur_position(&g);
      if (!mhash_insert(ht, sz, occ, position_get_hash_key(pos))) {
	game_cleanup(&g);
	return -1;
      }
    }
    if (len > 0)
      game_undo_move(&g);
    --len;
  }
  
  game_cleanup(&g);
  return ret;
}


int book_load(const char* filename)
{
  FILE* f = fopen(filename, "rb");
  if (f == 0) {
    error_set(ERR_IO);
    return 0;
  }

  uint8_t buf[16];
  if (fread(buf, 1, 8, f) != 8) {
    error_set(ERR_IO);
    fclose(f);
    return 0;
  }
  
  if (decode_uint64(buf) != BOOK_MAGIC) {
    error_set(ERR_BOOK_FORMAT);
    fclose(f);
    return 0;
  }
  if (fread(buf, 1, 8, f) != 8) {
    error_set(ERR_IO);
    fclose(f);
    return 0;
  }
  uint64_t n = decode_uint64(buf);
  uint64_t sz = 64;
  while (sz < (n * 100) / MHASH_LOAD_FACTOR)
    sz *= 2;
  
  hash_key_t* bh = (hash_key_t*)malloc(sizeof(hash_key_t) * sz);
  if (bh == 0) {
    error_set(ERR_MEMORY);
    fclose(f);
    return 0;
  }
  for (size_t i = 0; i < sz; i++)
    bh[i] = 0;

  for (size_t i = 0; i < n; i++) {
    if (fread(buf, 1, 16, f) != 16) {
      error_set(ERR_IO);
      free(bh);
      fclose(f);
      return 0;
    }
    hash_key_t k = decode_uint64(buf);
    uint64_t c = decode_uint64(buf + 8);
    hash_key_t index = k & (sz - 1);
    int inc = 1;
    while (bh[index] != 0 &&
	   BOOK_GET_KEY(bh[index]) != BOOK_GET_KEY(k)) {
      index = (index + inc)  & (sz - 1);
      ++inc;
    }
    c += BOOK_GET_COUNT(bh[index]);
    if (c > BOOK_MAX_COUNT)
      c = BOOK_MAX_COUNT;
    bh[index] = BOOK_ENTRY(k, c);
  }
  
  free(book_hash);
  book_hash = bh;
  hash_size = sz;
  fclose(f);
  return 1;
}


move_t book_move(struct game_t* g)
{
  assert(g != 0);

  if (book_hash == 0)
    return NULL_MOVE;

  move_t m[MAX_N_MOVES];
  const position_t* pos = game_cur_position(g);
  int n = movegen_legal_moves(pos, m);
  uint64_t tot = 0;
  uint64_t count[MAX_N_MOVES];
  for (int i = 0; i < n; i++) {
    game_make_move(g, m[i]);
    hash_key_t k = position_get_hash_key(pos);
    game_undo_move(g);
    
    hash_key_t index = k & (hash_size - 1);
    int inc = 1;
    while (book_hash[index] != 0 &&
	   BOOK_GET_KEY(book_hash[index]) != BOOK_GET_KEY(k)) {
      index = (index + inc)  & (hash_size - 1);
      ++inc;
    }
    if (book_hash[index] != 0) {
      count[i] = BOOK_GET_COUNT(book_hash[index]);
      tot += count[i];
    } else {
      count[i] = 0;
    }
  }

  if (tot == 0)
    return NULL_MOVE;

  uint64_t r = 0;
  for (int i = 0; i < 8; i++)
    r = (r << 7) + rand();
  r = r % tot;
  int sel = 0;
  while (sel < n - 1 && r >= count[sel]) {
    r -= count[sel];
    ++sel;
  }

  return m[sel];
}


int book_make(const char* pgnfile,
	      const char* bookfile,
	      int plies,
	      int mingames)
{
  assert(pgnfile != 0);
  assert(bookfile != 0);
  assert(plies >= 0);

  FILE* pgn = fopen(pgnfile, "rt");
  if (pgn == 0) {
    error_set(ERR_IO);
    return 0;
  }

  FILE* book = fopen(bookfile, "wb");
  if (book == 0) {
    error_set(ERR_IO);
    fclose(pgn);
    return 0;
  }
  
  size_t n = 0;
  struct mhash_entry_t* ht = 0;
  size_t occ = 0;
  
  int ret = 1;
  int ngames = 0;
  while (ret > 0) {
    ret = insert_pgn_game(pgn, &ht, &n, &occ, plies);
    ++ngames;
  }
  
  fclose(pgn);
  if (ret < 0) {
    free(ht);
    fclose(book);
    return 0;
  }
  
  size_t filtered = 0;
  for (size_t i = 0; i < n; i++)
    if (ht[i].count >= mingames)
      ++filtered;
  
  printf("%d games parsed: %zu unique positions will be stored.\n",
	 ngames, filtered);

  // Write the book.
  uint8_t buf[8];
  fwrite(encode_uint64(BOOK_MAGIC, buf), 1, 8, book);
  fwrite(encode_uint64(filtered, buf), 1, 8, book);
  for (size_t i = 0; i < n; i++) {
    if (ht[i].count < mingames)
      continue;
    if (fwrite(encode_uint64(ht[i].key, buf), 1, 8, book) != 8 ||
	fwrite(encode_uint64(ht[i].count, buf), 1, 8, book) != 8) {
      error_set(ERR_IO);
      ret = -1;
      break;
    }
  }
  
  free(ht);
  fclose(book);
  return (ret == 0 ? 1 : 0);
}


void book_cleanup()
{
  free(book_hash);
  book_hash = 0;
  hash_size = 0;
}

